home *** CD-ROM | disk | FTP | other *** search
/ Computer Shopper 242 / Issue 242 - April 2008 - DPCS0408DVD.ISO / Software Money Savers / VirtualDub / Source / VirtualDub-1.7.7-src.7z / src / system / source / thunk.cpp < prev    next >
Encoding:
C/C++ Source or Header  |  2007-09-05  |  7.7 KB  |  305 lines
  1. #include "stdafx.h"
  2. #include <windows.h>
  3. #include <map>
  4. #include <vd2/system/atomic.h>
  5. #include <vd2/system/refcount.h>
  6. #include <vd2/system/thunk.h>
  7. #include <vd2/system/binary.h>
  8.  
  9. class IVDJITAllocator {};
  10.  
  11. class VDJITAllocator : public vdrefcounted<IVDJITAllocator> {
  12. public:
  13.     VDJITAllocator();
  14.     ~VDJITAllocator();
  15.  
  16.     void *Allocate(size_t len);
  17.     void Free(void *p, size_t len);
  18.  
  19.     void EndUpdate(void *p, size_t len);
  20.  
  21. protected:
  22.     typedef std::map<void *, size_t> FreeChunks;
  23.     FreeChunks mFreeChunks;
  24.     FreeChunks::iterator mNextChunk;
  25.  
  26.     typedef std::map<void *, size_t> Allocations;
  27.     Allocations mAllocations;
  28.  
  29.     uintptr        mAllocationGranularity;
  30. };
  31.  
  32. VDJITAllocator::VDJITAllocator()
  33.     : mNextChunk(mFreeChunks.end())
  34. {
  35.     SYSTEM_INFO si;
  36.     GetSystemInfo(&si);
  37.  
  38.     mAllocationGranularity = si.dwAllocationGranularity;
  39. }
  40.  
  41. VDJITAllocator::~VDJITAllocator() {
  42.     for(Allocations::iterator it(mAllocations.begin()), itEnd(mAllocations.end()); it!=itEnd; ++it) {
  43.         VirtualFree(it->first, 0, MEM_RELEASE);
  44.     }
  45. }
  46.  
  47. void *VDJITAllocator::Allocate(size_t len) {
  48.     len = (len + 15) & ~(size_t)15;
  49.  
  50.     FreeChunks::iterator itMark(mNextChunk), itEnd(mFreeChunks.end()), it(itMark);
  51.  
  52.     if (it == itEnd)
  53.         it = mFreeChunks.begin();
  54.  
  55.     for(;;) {
  56.         for(; it!=itEnd; ++it) {
  57.             if (it->second >= len) {
  58.                 it->second -= len;
  59.  
  60.                 void *p = (char *)it->first + it->second;
  61.  
  62.                 if (!it->second) {
  63.                     if (mNextChunk == it)
  64.                         ++mNextChunk;
  65.  
  66.                     mFreeChunks.erase(it);
  67.                 }
  68.  
  69.                 return p;
  70.             }
  71.         }
  72.  
  73.         if (itEnd == itMark)
  74.             break;
  75.  
  76.         it = mFreeChunks.begin();
  77.         itEnd = itMark;
  78.     }
  79.  
  80.     size_t alloclen = (len + mAllocationGranularity - 1) & ~(mAllocationGranularity - 1);
  81.  
  82.     void *p = VirtualAlloc(NULL, alloclen, MEM_COMMIT, PAGE_EXECUTE_READWRITE);
  83.     if (p) {
  84.         try {
  85.             Allocations::iterator itA(mAllocations.insert(Allocations::value_type(p, alloclen)).first);
  86.  
  87.             try {
  88.                 if (len < alloclen)
  89.                     mFreeChunks.insert(FreeChunks::value_type((char *)p + len, alloclen - len));
  90.  
  91.             } catch(...) {
  92.                 mAllocations.erase(itA);
  93.                 throw;
  94.             }
  95.         } catch(...) {
  96.             VirtualFree(p, 0, MEM_RELEASE);
  97.             p = NULL;
  98.         }
  99.     }
  100.  
  101.     return p;
  102. }
  103.  
  104. void VDJITAllocator::Free(void *p, size_t len) {
  105.     FreeChunks::iterator cur(mFreeChunks.lower_bound(p));
  106.     if (cur != mFreeChunks.end() && (char *)p + len == cur->first) {
  107.         p = cur->first;
  108.         len += cur->second;
  109.         if (mNextChunk == cur)
  110.             ++mNextChunk;
  111.         cur = mFreeChunks.erase(cur);
  112.     }
  113.  
  114.     if (cur != mFreeChunks.begin()) {
  115.         FreeChunks::iterator prev(cur);
  116.  
  117.         --prev;
  118.         if ((char *)prev->first + prev->second == p) {
  119.             p = prev->first;
  120.             len += prev->second;
  121.             if (mNextChunk == prev)
  122.                 ++mNextChunk;
  123.             mFreeChunks.erase(prev);
  124.         }
  125.     }
  126.  
  127.     uintptr start = (size_t)p;
  128.     uintptr end = start + len;
  129.  
  130.     if (!((start | end) & (mAllocationGranularity - 1))) {
  131.         Allocations::iterator it(mAllocations.find(p));
  132.  
  133.         if (it != mAllocations.end()) {
  134.             VirtualFree((void *)start, 0, MEM_RELEASE);
  135.             mAllocations.erase(it);
  136.             return;
  137.         }
  138.     }
  139.  
  140.     mFreeChunks.insert(FreeChunks::value_type((void *)start, end-start));
  141. }
  142.  
  143. void VDJITAllocator::EndUpdate(void *p, size_t len) {
  144.     FlushInstructionCache(GetCurrentProcess(), p, len);
  145. }
  146.  
  147. ///////////////////////////////////////////////////////////////////////////
  148.  
  149. VDJITAllocator *g_pVDJITAllocator;
  150. VDAtomicInt g_VDJITAllocatorLock;
  151.  
  152. bool VDInitThunkAllocator() {
  153.     bool success = true;
  154.  
  155.     while(g_VDJITAllocatorLock.xchg(1))
  156.         ::Sleep(1);
  157.  
  158.     if (!g_pVDJITAllocator) {
  159.         g_pVDJITAllocator = new_nothrow VDJITAllocator;
  160.         if (!g_pVDJITAllocator)
  161.             success = false;
  162.     }
  163.  
  164.     if (success)
  165.         g_pVDJITAllocator->AddRef();
  166.  
  167.     VDVERIFY(1 == g_VDJITAllocatorLock.xchg(0));
  168.  
  169.     return success;
  170. }
  171.  
  172. void VDShutdownThunkAllocator() {
  173.     while(g_VDJITAllocatorLock.xchg(1))
  174.         ::Sleep(1);
  175.  
  176.     VDASSERT(g_pVDJITAllocator);
  177.  
  178.     if (!g_pVDJITAllocator->Release())
  179.         g_pVDJITAllocator = NULL;
  180.  
  181.     VDVERIFY(1 == g_VDJITAllocatorLock.xchg(0));
  182. }
  183.  
  184. void *VDAllocateThunkMemory(size_t len) {
  185.     return g_pVDJITAllocator->Allocate(len);
  186. }
  187.  
  188. void VDFreeThunkMemory(void *p, size_t len) {
  189.     g_pVDJITAllocator->Free(p, len);
  190. }
  191.  
  192. void VDSetThunkMemory(void *p, const void *src, size_t len) {
  193.     memcpy(p, src, len);
  194.     g_pVDJITAllocator->EndUpdate(p, len);
  195. }
  196.  
  197. void VDFlushThunkMemory(void *p, size_t len) {
  198.     g_pVDJITAllocator->EndUpdate(p, len);
  199. }
  200.  
  201. ///////////////////////////////////////////////////////////////////////////
  202.  
  203. #ifdef _M_AMD64
  204.     extern "C" void VDMethodToFunctionThunk64();
  205. #else
  206.     extern "C" void VDMethodToFunctionThunk32();
  207.     extern "C" void VDMethodToFunctionThunk32_4();
  208.     extern "C" void VDMethodToFunctionThunk32_8();
  209.     extern "C" void VDMethodToFunctionThunk32_12();
  210.     extern "C" void VDMethodToFunctionThunk32_16();
  211. #endif
  212.  
  213. VDFunctionThunk *VDCreateFunctionThunkFromMethod(void *method, void *pThis, size_t argbytes, bool stdcall_thunk) {
  214. #if defined(_M_IX86)
  215.     void *pThunk = VDAllocateThunkMemory(16);
  216.  
  217.     if (!pThunk)
  218.         return NULL;
  219.  
  220.     if (stdcall_thunk || !argbytes) {    // thiscall -> stdcall (easy case)
  221.         uint8 thunkbytes[16]={
  222.             0xB9, 0x00, 0x00, 0x00, 0x00,                // mov ecx, this
  223.             0xE9, 0x00, 0x00, 0x00, 0x00                // jmp fn
  224.         };
  225.  
  226.  
  227.         VDWriteUnalignedLEU32(thunkbytes+1, (uint32)(uintptr)pThis);
  228.         VDWriteUnalignedLEU32(thunkbytes+6, (uint32)method - ((uint32)pThunk + 10));
  229.  
  230.         VDSetThunkMemory(pThunk, thunkbytes, 15);
  231.     } else {                // thiscall -> cdecl (hard case)
  232.         uint8 thunkbytes[16]={
  233.             0xE8, 0x00, 0x00, 0x00, 0x00,                // call VDFunctionThunk32
  234.             0xC3,                                        // ret
  235.             argbytes,                                    // db argbytes
  236.             0,                                            // db 0
  237.             0x00, 0x00, 0x00, 0x00,                        // dd method
  238.             0x00, 0x00, 0x00, 0x00,                        // dd this
  239.         };
  240.  
  241.         void *adapter;
  242.  
  243.         switch(argbytes) {
  244.         case 4:        adapter = VDMethodToFunctionThunk32_4;    break;
  245.         case 8:        adapter = VDMethodToFunctionThunk32_8;    break;
  246.         case 12:    adapter = VDMethodToFunctionThunk32_12;    break;
  247.         case 16:    adapter = VDMethodToFunctionThunk32_16;    break;
  248.         default:    adapter = VDMethodToFunctionThunk32;    break;
  249.         }
  250.  
  251.         VDWriteUnalignedLEU32(thunkbytes+1, (uint32)(uintptr)adapter - ((uint32)pThunk + 5));
  252.         VDWriteUnalignedLEU32(thunkbytes+8, (uint32)(uintptr)method);
  253.         VDWriteUnalignedLEU32(thunkbytes+12, (uint32)(uintptr)pThis);
  254.  
  255.         VDSetThunkMemory(pThunk, thunkbytes, 16);
  256.     }
  257.  
  258.     return (VDFunctionThunk *)pThunk;
  259. #elif defined(_M_AMD64)
  260.     void *pThunk = VDAllocateThunkMemory(44);
  261.     if (!pThunk)
  262.         return NULL;
  263.  
  264.     uint8 thunkbytes[44]={
  265.         0x48, 0x8D, 0x04, 0x25, 0x10, 0x00, 0x00,    // lea rax, [eip+16]
  266.         0x00,
  267.         0xFF, 0x24, 0x25, 0x08, 0x00, 0x00, 0x00,    // jmp qword ptr [rip+8]
  268.         0x90,                                        // nop
  269.         0, 0, 0, 0, 0, 0, 0, 0,                        // dq VDFunctionThunk64
  270.         0, 0, 0, 0, 0, 0, 0, 0,                        // dq method
  271.         0, 0, 0, 0, 0, 0, 0, 0,                        // dq this
  272.         0, 0, 0, 0                                    // dd argspillbytes
  273.     };
  274.  
  275.     VDWriteUnalignedLEU64(thunkbytes+16, (uint64)(uintptr)VDMethodToFunctionThunk64);
  276.     VDWriteUnalignedLEU64(thunkbytes+24, (uint64)(uintptr)method);
  277.     VDWriteUnalignedLEU64(thunkbytes+32, (uint64)(uintptr)pThis);
  278.  
  279.     // The stack must be aligned to a 16 byte boundary when the CALL
  280.     // instruction occurs. On entry to VDFunctionThunk64(), the stack is misaligned
  281.     // to 16n+8. Therefore, the number of argbytes must be 16m+8 and the number of
  282.     // argspillbytes must be 16m+8-24.
  283.     VDWriteUnalignedLEU32(thunkbytes+40, argbytes < 32 ? 0 : ((argbytes - 16 + 15) & ~15));
  284.  
  285.     VDSetThunkMemory(pThunk, thunkbytes, 44);
  286.  
  287.     return (VDFunctionThunk *)pThunk;
  288. #else
  289.     return NULL;
  290. #endif
  291. }
  292.  
  293. void VDDestroyFunctionThunk(VDFunctionThunk *pFnThunk) {
  294.     // validate thunk
  295. #if defined(_M_IX86)
  296.     VDASSERT(((const uint8 *)pFnThunk)[0] == 0xB9 || ((const uint8 *)pFnThunk)[0] == 0xE8);
  297.     VDFreeThunkMemory(pFnThunk, 16);
  298. #elif defined(_M_AMD64)
  299.     VDFreeThunkMemory(pFnThunk, 44);
  300. #else
  301.     VDASSERT(false);
  302. #endif
  303.  
  304. }
  305.